# -*- coding: utf-8 -*-
"""
训练集：train_aligned.csv
测试集：test_final.csv
模型：LightGBM + Top20 特征
输出：测试集 AUC、分类报告、混淆矩阵
"""

import pandas as pd
import numpy as np
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.metrics import roc_auc_score, classification_report, confusion_matrix
import lightgbm as lgb

# 1. 读数据
train = pd.read_csv('.\data\processed\\train_final.csv')
test  = pd.read_csv('.\data\processed\\test_final.csv')

target = 'Attrition'
y_train = train[target].astype(int)
y_test  = test[target].astype(int)

X_train = train.drop(columns=[target])
X_test  = test.drop(columns=[target])

# 2. 预处理（与训练时完全一致）
num_cols = X_train.select_dtypes(include=['int64','float64']).columns
cat_cols = X_train.select_dtypes(include=['object','category','bool']).columns

pre = ColumnTransformer(
        [('num', StandardScaler(), num_cols),
         ('cat', OneHotEncoder(handle_unknown='ignore'), cat_cols)])

X_train_enc = pre.fit_transform(X_train)
X_test_enc  = pre.transform(X_test)

# 3. 特征选择 – 用训练集训练一次 LightGBM 并取 Top20
lgb0 = lgb.LGBMClassifier(n_estimators=100, random_state=42, verbose=-1)
lgb0.fit(X_train_enc, y_train)
imp = lgb0.feature_importances_
top20 = np.argsort(imp)[-20:]

X_train_top = X_train_enc[:, top20]
X_test_top  = X_test_enc[:, top20]

# 4. 最终模型（与复现相同的超参）
lgb_final = lgb.LGBMClassifier(
        n_estimators=400,
        learning_rate=0.05,
        subsample=0.8,
        colsample_bytree=0.8,
        random_state=42,
        verbose=-1)

lgb_final.fit(X_train_top, y_train)

# 5. 预测与评估
y_pred_prob = lgb_final.predict_proba(X_test_top)[:, 1]
y_pred = (y_pred_prob >= 0.5).astype(int)

print('=== 测试集结果 ===')
print('AUC :', roc_auc_score(y_test, y_pred_prob))
print('\n分类报告：')
print(classification_report(y_test, y_pred, digits=4))
print('混淆矩阵：')
print(confusion_matrix(y_test, y_pred))